from fastapi import FastAPI, HTTPException
from contextlib import asynccontextmanager
import chromadb
import os
from HydrogenExpertSystem import model


@asynccontextmanager
async def lifespan(app: FastAPI):
    client = chromadb.PersistentClient(path=os.path.abspath("./chromadb"))
    app.state.collection = client.get_or_create_collection(
        name="hydrogen_research",
        metadata={"hnsw:space": "cosine"}
    )
    print(f"初始化完成，记录数：{app.state.collection.count()}")
    yield

app = FastAPI(lifespan=lifespan)

@app.post("/search")
async def search(request: dict):
    try:
        text = request["text"]
        embedding = model.encode([text]).tolist()[0]
        results = app.state.collection.query(
            query_embeddings=[embedding],
            n_results=3
        )
        return {
            "documents": results["documents"][0],
            "scores": [1 - score for score in results["distances"][0]]
        }
    except KeyError as e:
        raise HTTPException(400, f"缺少必要参数: {str(e)}")

@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

